import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

class RerankerModel:
    def __init__(self, model_name, device='cuda'):
        """
        Initializes the reranker model with the specified model name and device.
        """
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True)
        self.device = device if torch.cuda.is_available() else 'cpu'
        self.model.to(self.device)

    def score_pairs(self, sentence_pairs):
        """
        Processes a list of sentence pairs to produce their scores.
        """
        inputs = self.tokenizer(sentence_pairs, padding=True, truncation=True, max_length=512, return_tensors="pt")
        inputs_on_device = {k: v.to(self.device) for k, v in inputs.items()}
        logits = self.model(**inputs_on_device, return_dict=True).logits.view(-1,).float()
        scores = torch.sigmoid(logits)
        return scores

if __name__ == '__main__':
    # Usage example:
    model_name = '/home/jina-reranker-v2-base-multilingual'
    sentence_pairs = [
        ("The weather is nice today.", "It's a beautiful day."),
        ("He won the race.", "He came first in the competition.")
    ]
    reranker_model = RerankerModel(model_name, device='cuda')
    scores = reranker_model.score_pairs(sentence_pairs)
    print(scores)
